package net.lulihu.pangolin.headless.common.util;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.pool.KryoPool;
import net.lulihu.pangolin.headless.common.protocol.ProtocolMessage;
import org.objenesis.strategy.StdInstantiatorStrategy;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.util.function.Function;

/**
 * Kryo 消息协议序列化工具
 */
public class KryoProtocolMessageKit {

    /**
     * KryoPool 消息协议线程池
     */
    private static KryoPool kryoPool = new KryoPool.Builder(() -> {
        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.setRegistrationRequired(false);
        kryo.register(ProtocolMessage.class);
        //在反序列化集合时修复NPE错误.
        ((Kryo.DefaultInstantiatorStrategy) kryo.getInstantiatorStrategy())
                .setFallbackInstantiatorStrategy(new StdInstantiatorStrategy());
        return kryo;
    }).build();

    /**
     * 使用kryo对象执行操作
     */
    private static <T> T kryo(Function<Kryo, T> function) {
        Kryo kryo = kryoPool.borrow();
        try {
            return function.apply(kryo);
        } finally {
            kryoPool.release(kryo);
        }
    }

    //-----------------------------------------------
    //          序列化/反序列化对象，及类型信息
    //          序列化的结果里，包含类型的信息
    //          反序列化时不再需要提供类型
    //-----------------------------------------------

    /**
     * 将对象【及类型】序列化为字节数组
     *
     * @param obj 任意对象
     * @return 序列化后的字节数组
     */
    public static byte[] writeToByteArray(final Object obj) {
        return kryo(kryo -> {
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            try (Output output = new Output(byteArrayOutputStream)) {
                kryo.writeClassAndObject(output, obj);
                return output.toBytes();
            }
        });
    }


    /**
     * 将字节数组反序列化为原对象
     *
     * @param byteArray writeToByteArray 方法序列化后的字节数组
     * @param <T>       原对象的类型
     * @return 原对象
     */
    @SuppressWarnings("unchecked")
    public static <T> T readFromByteArray(byte[] byteArray) {
        return kryo(kryo -> {
            ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(byteArray);
            try (Input input = new Input(byteArrayInputStream)) {
                return (T) kryo.readClassAndObject(input);
            }
        });
    }

    //-----------------------------------------------
    //          只序列化/反序列化对象
    //          序列化的结果里，不包含类型的信息
    //-----------------------------------------------

    /**
     * 将对象序列化为字节数组
     *
     * @param obj 任意对象
     * @return 序列化后的字节数组
     */
    public static byte[] writeObjectToByteArray(Object obj) {
        return kryo(kryo -> {
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            try (Output output = new Output(byteArrayOutputStream)) {
                kryo.writeObject(output, obj);
                return output.toBytes();
            }
        });
    }

    /**
     * 将字节数组反序列化为原对象
     *
     * @param byteArray writeObjectToByteArray 方法序列化后的字节数组
     * @param clazz     原对象的 Class
     * @param <T>       原对象的类型
     * @return 原对象
     */
    public static <T> T readObjectFromByteArray(byte[] byteArray, Class<T> clazz) {
        return kryo(kryo -> {
            ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(byteArray);
            try (Input input = new Input(byteArrayInputStream)) {
                return kryo.readObject(input, clazz);
            }
        });
    }


    /**
     * 将字节数组反序列化为原对象
     *
     * @param input 输入流
     * @param clazz 原对象的 Class
     * @param <T>   原对象的类型
     * @return 原对象
     */
    public static <T> T readObjectFromInput(Input input, Class<T> clazz) {
        return kryo(kryo -> {
            try {
                return kryo.readObject(input, clazz);
            } finally {
                input.close();
            }
        });
    }

}